package tech.waterism.model.util;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.serializer.SerializerFeature;
import tech.waterism.model.dto.ModelOutputSimple;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.*;

/**
 * 模型输出辅助工具
 * 
 * <p>用于构建扁平化的模型输出并保存到JSON文件</p>
 * 
 * @author AI Assistant
 * @version 1.0
 * @since 2025-11-10
 */
public class ModelOutputHelper {
    
    private static final SimpleDateFormat DATE_FORMAT = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
    
    private final String modelName;
    private final String modelVersion;
    private Date[] timeStamps;
    private Integer warmUpPeriods = 0;
    private final Map<String, double[]> dataArrays = new LinkedHashMap<>();
    private final Map<String, double[]> stateArrays = new LinkedHashMap<>();
    private Map<String, Object> finalState;
    private Map<String, Double> modelParameters;
    private Map<String, Object> metadata;
    
    public ModelOutputHelper(String modelName, String modelVersion) {
        this.modelName = modelName;
        this.modelVersion = modelVersion;
    }
    
    public ModelOutputHelper setTimeStamps(Date[] timeStamps) {
        this.timeStamps = timeStamps;
        return this;
    }
    
    public ModelOutputHelper setWarmUpPeriods(int warmUpPeriods) {
        this.warmUpPeriods = warmUpPeriods;
        return this;
    }
    
    public ModelOutputHelper addDataArray(String key, double[] values) {
        if (values != null) {
            this.dataArrays.put(key, values);
        }
        return this;
    }
    
    /**
     * 添加状态变量数组
     * 
     * @param key 状态变量名称
     * @param values 状态变量值数组
     * @return this
     */
    public ModelOutputHelper addStateArray(String key, double[] values) {
        if (values != null) {
            this.stateArrays.put(key, values);
        }
        return this;
    }
    
    public ModelOutputHelper setFinalState(Map<String, Object> finalState) {
        this.finalState = finalState;
        return this;
    }
    
    public ModelOutputHelper setModelParameters(Map<String, Double> modelParameters) {
        this.modelParameters = modelParameters;
        return this;
    }
    
    public ModelOutputHelper setMetadata(Map<String, Object> metadata) {
        this.metadata = metadata;
        return this;
    }
    
    /**
     * 构建扁平化输出
     */
    public ModelOutputSimple build() {
        if (timeStamps == null || timeStamps.length == 0) {
            throw new IllegalStateException("时间序列不能为空");
        }
        
        int totalPeriods = timeStamps.length;
        
        // 构建时间字符串数组
        String[] timeStrings = new String[totalPeriods];
        Boolean[] isWarmUpArray = new Boolean[totalPeriods];
        for (int i = 0; i < totalPeriods; i++) {
            timeStrings[i] = DATE_FORMAT.format(timeStamps[i]);
            isWarmUpArray[i] = i < warmUpPeriods;
        }
        
        // 计算起报时间
        String forecastStartTime = warmUpPeriods < totalPeriods ? 
                timeStrings[warmUpPeriods] : timeStrings[totalPeriods - 1];
        
        // 计算预报期时段数
        int forecastPeriods = totalPeriods - warmUpPeriods;
        
        // 构建时间序列数组：每个元素包含时间、isWarmUp和所有数据字段
        List<Map<String, Object>> timeSeriesList = new ArrayList<>();
        for (int i = 0; i < totalPeriods; i++) {
            Map<String, Object> dataPoint = new LinkedHashMap<>();
            dataPoint.put("time", timeStrings[i]);
            dataPoint.put("isWarmUp", isWarmUpArray[i]);
            
            // 添加所有数据字段
            for (Map.Entry<String, double[]> entry : dataArrays.entrySet()) {
                double[] values = entry.getValue();
                if (values != null && i < values.length) {
                    dataPoint.put(entry.getKey(), values[i]);
                }
            }
            
            timeSeriesList.add(dataPoint);
        }
        
        // 构建状态变量时间序列
        List<Map<String, Object>> stateTimeSeriesList = null;
        if (!stateArrays.isEmpty()) {
            stateTimeSeriesList = new ArrayList<>();
            // 状态数组的长度可能是 totalPeriods + 1（包含初始状态）
            int stateLength = stateArrays.values().iterator().next().length;
            
            for (int i = 0; i < stateLength; i++) {
                Map<String, Object> statePoint = new LinkedHashMap<>();
                
                // 添加时间戳（如果索引超出时间数组，使用最后一个时间）
                if (i < totalPeriods) {
                    statePoint.put("time", timeStrings[i]);
                    statePoint.put("isWarmUp", isWarmUpArray[i]);
                } else {
                    // 最后一个状态点（对应最后一个时段结束时的状态）
                    statePoint.put("time", timeStrings[totalPeriods - 1] + " (end)");
                    statePoint.put("isWarmUp", false);
                }
                
                // 添加所有状态变量
                for (Map.Entry<String, double[]> entry : stateArrays.entrySet()) {
                    double[] values = entry.getValue();
                    if (values != null && i < values.length) {
                        statePoint.put(entry.getKey(), values[i]);
                    }
                }
                
                stateTimeSeriesList.add(statePoint);
            }
        }
        
        // 构建输出对象
        ModelOutputSimple output = new ModelOutputSimple();
        output.setModelName(modelName);
        output.setModelVersion(modelVersion);
        output.setComputeTime(DATE_FORMAT.format(new Date()));
        output.setWarmUpPeriods(warmUpPeriods);
        output.setForecastPeriods(forecastPeriods);
        output.setForecastStartTime(forecastStartTime);
        output.setTimeSeries(timeSeriesList);
        output.setStateTimeSeries(stateTimeSeriesList);
        output.setFinalState(finalState);
        output.setModelParameters(modelParameters);
        output.setMetadata(metadata);
        
        return output;
    }
    
    /**
     * 保存输出到JSON文件
     * 
     * @param output 输出对象
     * @param filePath 文件路径
     * @throws IOException IO异常
     */
    public static void saveToFile(ModelOutputSimple output, String filePath) throws IOException {
        File file = new File(filePath);
        File parentDir = file.getParentFile();
        if (parentDir != null && !parentDir.exists()) {
            parentDir.mkdirs();
        }
        
        try (FileWriter writer = new FileWriter(file)) {
            String json = JSON.toJSONString(output, 
                    SerializerFeature.PrettyFormat,
                    SerializerFeature.WriteMapNullValue,
                    SerializerFeature.DisableCircularReferenceDetect);
            writer.write(json);
        }
    }
}

